import numpy as np
import PIL.Image as Image
import patchreconst as pr
import matplotlib.pyplot as plt

PATCH_SIZE = 32
COMPONENT_COUNT = 350
PATCH_COUNT = 400000
EPOCH_COUNT = 4
CODE_COUNT = 2800
SPARSITY_PARAMETER = 2.0
SPARSITY_PARAMETER2 = 4.0
NON_NEGATIVE = True

scBasisFile = 'basis/basisnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'
scBasis = np.load(scBasisFile)

scBasisFile2 = 'basis/basisnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER2) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'
scBasis2 = np.load(scBasisFile2)

icaFilterFile = 'icabasis/basisluisicacomps' + str(COMPONENT_COUNT) + 'codes' + str(CODE_COUNT) + 'patches' + str(PATCH_COUNT) + '.npy'
icaFilters = np.load(icaFilterFile)

patches = np.load('patches_1000_mod.npy')
patchesOrig = np.load('patches_1000.npy')[:, :, :, 0]
rmBounds = np.load('patches_1000_rmbounds.npy')
rmBounds = (rmBounds[0], rmBounds[1], rmBounds[2])

scCodes, icaCodes, v1Simple, v1cMean, angles = pr.responses(patches, scBasis, SPARSITY_PARAMETER, NON_NEGATIVE, icaFilters)
scCodes2, _, _, _, _ = pr.responses(patches, scBasis2, SPARSITY_PARAMETER2, NON_NEGATIVE, icaFilters)
pcaTransformed, v1C = pr.responsesPCAV1C(patches)

#angles = np.random.normal(size = angles.shape)
#angles = np.random.vonmises(-np.pi / 6.0, 2.0, size = angles.shape)
#angles = np.random.uniform(low = -np.pi, high = np.pi, size = angles.shape)
#np.save('angles.npy', angles)

patchesReconstV1 = pr.reconstructV1(v1Simple)
patchesReconstV1C = pr.reconstructV1C(v1C, v1cMean, angles)
patchesReconstPCA = pr.reconstructPCA(pcaTransformed, v1cMean, angles)
patchesReconstSC = pr.reconstruct(scBasis, scCodes, v1cMean, angles)
patchesReconstSC2 = pr.reconstruct(scBasis2, scCodes2, v1cMean, angles)
patchesReconstICA = pr.reconstruct(icaFilters, icaCodes, v1cMean, angles)

sc1error = np.mean((patchesReconstSC - patchesReconstV1).reshape((1000, -1)) ** 2.0, axis = -1)

icaerror = np.mean((patchesReconstICA - patchesReconstV1).reshape((1000, -1)) ** 2.0, axis = -1)

v1vsorig = np.mean((patchesReconstV1 - patchesOrig).reshape((1000, -1)) ** 2.0, axis = -1)

pcaError = np.mean((patchesReconstPCA - patchesReconstV1).reshape((1000, -1)) ** 2.0, axis = -1)

sc2error = np.mean((patchesReconstSC2 - patchesReconstV1).reshape((1000, -1)) ** 2.0, axis = -1)

#Uncomment to show patch representation at each stage of the model.
'''
for i in range(patchesReconstV1.shape[0]):
	f, (ax1, ax2, ax3, ax4, ax5, ax6, ax7) = plt.subplots(1, 7, sharey = True)

	ax1.imshow(patchesOrig[i], cmap = 'gray', interpolation = 'none')
#	ax1.title.set_text('Image')
	ax1.set_xticks([])
	ax1.set_yticks([])

	ax2.imshow(patchesReconstV1[i], cmap = 'gray', interpolation = 'none')
#	ax2.title.set_text('V1')
	ax2.set_xticks([])
	ax2.set_yticks([])

	ax3.imshow(patchesReconstV1C[i], cmap = 'gray', interpolation = 'none')
#	ax3.title.set_text('V1C Mod')
	ax3.set_xticks([])
	ax3.set_yticks([])

	ax4.imshow(patchesReconstPCA[i], cmap = 'gray', interpolation = 'none')
#	ax4.title.set_text('PCA')
	ax4.set_xticks([])
	ax4.set_yticks([])

	ax5.imshow(patchesReconstICA[i], cmap = 'gray', interpolation = 'none')
#	ax5.title.set_text('ICA')
	ax5.set_xticks([])
	ax5.set_yticks([])

	ax6.imshow(patchesReconstSC[i], cmap = 'gray', interpolation = 'none')
#	ax6.title.set_text('SC ' + str(SPARSITY_PARAMETER))
	ax6.set_xticks([])
	ax6.set_yticks([])

	ax7.imshow(patchesReconstSC2[i], cmap = 'gray', interpolation = 'none')
#	ax7.title.set_text('SC ' + str(SPARSITY_PARAMETER2))
	ax7.set_xticks([])
	ax7.set_yticks([])

	plt.savefig('pred/' + str(i) + '.png')
	plt.close()
'''
